{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analyzing the gradients of the SHT\n",
    "\n",
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import paddle\n",
    "import paddle.nn as nn\n",
    "\n",
    "from paddle_harmonics.quadrature import legendre_gauss_weights, clenshaw_curtiss_weights\n",
    "from paddle_harmonics.legendre import legpoly, clm\n",
    "from paddle_harmonics import RealSHT, InverseRealSHT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem setting\n",
    "\n",
    "We consider the simple problem of fitting the spectral coefficients $\\theta$ such that\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\mathcal{L}\n",
    "&= ||\\mathop{\\mathrm{ISHT}}[\\theta] - u^*||^2_{S^2} \\\\\n",
    "&\\approx \\sum_j \\omega_j (\\mathop{\\mathrm{ISHT}}[\\theta](x_j) - u^*(x_j))^2 \\\\\n",
    "&= (S \\, \\theta - u^*)^T \\mathop{\\mathrm{diag}}(\\omega) \\, (S \\, \\theta - u^*) \\\\\n",
    "&= L\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "is minimized.\n",
    "\n",
    "The Vandermonde matrix $S$, which is characterized by $\\mathop{\\mathrm{ISHT}}[\\theta] = S \\theta$ realizes the action of the discrete SHT.\n",
    "\n",
    "The necessary condition for a minimizer of $L$ is\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "& \\nabla_\\theta L = S^T \\mathop{\\mathrm{diag}}(\\omega) \\, (S \\, \\theta - u^*) = 0 \\\\\n",
    "\\Leftrightarrow \\quad & S^T \\mathop{\\mathrm{diag}}(\\omega) \\, S \\; \\theta = S^T \\mathop{\\mathrm{diag}}(\\omega) \\, u^*.\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "On the Gaussian grid, "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlat = 64\n",
    "nlon = 2*nlat\n",
    "grid = \"equiangular\"\n",
    "\n",
    "# for quadrature and plotting\n",
    "if grid == \"legendre-gauss\":\n",
    "    lmax = mmax = nlat\n",
    "    xq, wq = legendre_gauss_weights(nlat)\n",
    "elif grid ==\"equiangular\":\n",
    "    lmax = mmax = nlat//2\n",
    "    xq, wq = clenshaw_curtiss_weights(nlat)\n",
    "\n",
    "sht = RealSHT(nlat, nlon, lmax=lmax, mmax=mmax, grid=grid)\n",
    "isht = InverseRealSHT(nlat, nlon, lmax=lmax, mmax=mmax, grid=grid)\n",
    "\n",
    "lat = np.arccos(xq)\n",
    "omega = math.pi * paddle.to_tensor(wq).to(paddle.float32) / nlat\n",
    "omega = omega.reshape([-1, 1])\n",
    "\n",
    "nlon*omega.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p ./data\n",
    "!wget https://astropedia.astrogeology.usgs.gov/download/Mars/GlobalSurveyor/MOLA/thumbs/Mars_MGS_MOLA_DEM_mosaic_global_1024.jpg -O ./data/mola_topo.jpg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import imageio.v3 as iio\n",
    "\n",
    "img = iio.imread('./data/mola_topo.jpg')\n",
    "#convert to grayscale\n",
    "data = np.dot(img[...,:3]/255, [0.299, 0.587, 0.114])\n",
    "# interpolate onto 512x1024 grid:\n",
    "data = nn.functional.interpolate(paddle.to_tensor(data).unsqueeze(0).unsqueeze(0), size=(nlat,nlon)).squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from plotting import plot_sphere\n",
    "\n",
    "plot_sphere(data, cmap=\"turbo\", colorbar=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = 1.0\n",
    "theta = paddle.create_parameter([lmax,lmax,2], paddle.float32, \n",
    "                                default_initializer=paddle.nn.initializer.Normal(mean=0, std=1))\n",
    "optim = paddle.optimizer.SGD(parameters=[theta], learning_rate=lr)\n",
    "\n",
    "for iter in range(40):\n",
    "    optim.clear_grad(set_to_zero=True)\n",
    "    loss = paddle.sum(0.5*omega*(isht(theta.as_complex()) - data)**2)\n",
    "    loss.backward()\n",
    "\n",
    "    # action of the Hessian\n",
    "    with paddle.no_grad():\n",
    "        for m in range(1,mmax):\n",
    "            theta.grad[:,m].multiply_(paddle.to_tensor([0.5],dtype=paddle.float32))\n",
    "\n",
    "    optim.step()\n",
    "\n",
    "    print(f\"iter: {iter}, loss: {loss}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "what's the best possible loss? $\\theta^* = (S^T \\mathop{\\mathrm{diag}}(\\omega) \\, S)^{-1} S^T \\mathop{\\mathrm{diag}}(\\omega) u^* = \\mathop{\\mathrm{SHT}}[u^*]$ gives us the global minimizer for this problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(layout='constrained', figsize=(20, 12))\n",
    "subfigs = fig.subfigures(2, 3)\n",
    "\n",
    "# spectral fitting\n",
    "plot_sphere(isht(theta.as_complex()).detach(), fig=subfigs[0,0], cmap=\"turbo\", colorbar=True, title=\"Fit\")\n",
    "plot_sphere(data, fig=subfigs[0,1], cmap=\"turbo\", colorbar=True, title=\"Ground truth\")\n",
    "plot_sphere((isht(theta.as_complex()) - data).detach(), fig=subfigs[0,2], cmap=\"turbo\", colorbar=True, title=\"residual\")\n",
    "\n",
    "# sht(u)\n",
    "plot_sphere(isht(sht(data)).detach(), fig=subfigs[1,0], cmap=\"turbo\", colorbar=True, title=\"isht(sht(u))\")\n",
    "plot_sphere(data, fig=subfigs[1,1], cmap=\"turbo\", colorbar=True, title=\"Ground truth\")\n",
    "plot_sphere((isht(sht(data)) - data).detach(), fig=subfigs[1,2], cmap=\"turbo\", colorbar=True, title=\"residual\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
